from typing import Callable
from learners.baseline import Learner as BaselineLearner
from models.multi_net import BaseMulti

num_workers = 8


class Learner(BaselineLearner):
    _network: BaseMulti

    def __init__(self, args, model_func: Callable = BaseMulti):
        super().__init__(args, model_func)

        self.forward_train = lambda model, inputs: model(inputs, mode="cur")
        self.forward_test = lambda model, inputs: model(inputs, mode="all")

    def extract_token(self, inputs):
        """Only work with mode CUR, 兼容DataParallel"""
        return self._call_network_method("extract_token", inputs, mode="cur")[:, -1]
